{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "6601d262-aac9-48ba-9296-04d8c5f9966a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import sys\n",
    "\n",
    "import numpy as np\n",
    "import tritonclient.grpc as grpcclient\n",
    "import json\n",
    "\n",
    "import cv2\n",
    "import time\n",
    "from tqdm import tqdm\n",
    "import glob\n",
    "import os\n",
    "import os.path as osp\n",
    "from concurrent.futures import ThreadPoolExecutor\n",
    "import fiftyone as fo\n",
    "\n",
    "from my_class import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "11cf35d2-39b4-45e9-86e6-4af53aebebae",
   "metadata": {},
   "outputs": [],
   "source": [
    "url = \"0.0.0.0:8030\"\n",
    "triton_client = grpcclient.InferenceServerClient(\n",
    "    url=url,\n",
    "    verbose=False,\n",
    ")\n",
    "\n",
    "dataset_name=\"nike_infer\"\n",
    "fo.delete_datasets('*')\n",
    "\n",
    "dataset = fo.Dataset(name=dataset_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "bfa96e91-3533-45bb-bc71-d83d281d13de",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nums of batch:  1\n"
     ]
    }
   ],
   "source": [
    "metric_num = \"test0425exp1\"\n",
    "root_dir = f\"/home/ray/input/AI_Image/202504/{metric_num}/\"\n",
    "# root_dir = \"/home/user/workspace/datasets/tmp/res\"\n",
    "batch_nums = os.listdir(root_dir)\n",
    "\n",
    "\n",
    "# 随机10个批次\n",
    "# batch_nums = random.sample(batch_nums, 2)\n",
    "print(\"nums of batch: \", len(batch_nums))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "091b7baa-4ac8-4f6a-8f48-36dd24ac04b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = []\n",
    "def infer(image_data):\n",
    "    inputs = []\n",
    "    outputs = []\n",
    "    model_name = \"ensemble\"\n",
    "    input_name = \"image\"\n",
    "    skill_out_json = 'skill_out_json'\n",
    "    output_name = {skill_out_json: 0}\n",
    "    \n",
    "    for n in output_name.keys():\n",
    "        outputs.append(grpcclient.InferRequestedOutput(n))\n",
    "\n",
    "    # 1. Preprocess 时间\n",
    "    t0 = time.time()\n",
    "    inputs.append(grpcclient.InferInput(input_name, image_data.shape, \"UINT8\"))\n",
    "    inputs[0].set_data_from_numpy(image_data)\n",
    "    t1 = time.time()\n",
    "\n",
    "    # 2. 推理 + 网络调用时间\n",
    "    results = triton_client.infer(\n",
    "        model_name=model_name,\n",
    "        inputs=inputs,\n",
    "        request_id=str(\"Camrea1\"),\n",
    "        outputs=outputs\n",
    "    )\n",
    "    t2 = time.time()\n",
    "\n",
    "    # 3. Postprocess 时间\n",
    "    json_str = results.as_numpy(skill_out_json)\n",
    "    json_str = json_str.tobytes()\n",
    "    res = json.loads(json_str)\n",
    "    t3 = time.time()\n",
    "\n",
    "    # print(f\"[TIMING] Preprocess: {(t1 - t0)*1000:.2f} ms\")\n",
    "    # print(f\"[TIMING] Infer+Network: {(t2 - t1)*1000:.2f} ms\")\n",
    "    # print(f\"[TIMING] Postprocess: {(t3 - t2)*1000:.2f} ms\")\n",
    "    # print(f\"[TIMING] Total: {(t3 - t0)*1000:.2f} ms\")\n",
    "\n",
    "    return res\n",
    "\n",
    "def process_image(im_file):\n",
    "    sample = fo.Sample(filepath=im_file)\n",
    "    img_shape = (4096, 3000)\n",
    "    img_w = img_shape[0]\n",
    "    img_h = img_shape[1]\n",
    "    img = np.fromfile(im_file, dtype='uint8')\n",
    "    image_data = np.array(img)\n",
    "    image_data = np.expand_dims(image_data, axis=0)\n",
    "\n",
    "    id = \"Camera01\"\n",
    "    \n",
    "    res = infer(image_data)  \n",
    "    # 鞋子坐标\n",
    "    shoe = BBox(*res[0][\"offset\"][:4]) / (img_w, img_h) \n",
    "    detections = [fo.Detection(label=\"shoe\", bounding_box=shoe.xywh)]\n",
    "    sample[\"shoe\"] = fo.Detections(detections=detections)\n",
    "    \n",
    "\n",
    "    # 缺陷坐标\n",
    "    detections = []\n",
    "    segmentations = []\n",
    "    predicitons = res[0][\"predictions\"]\n",
    "    for flaw in predicitons:\n",
    "        item = Anomaly(flaw[\"scores\"], flaw[\"bboxes\"], flaw[\"segmentation\"], flaw[\"category_name\"])\n",
    "        item.bbox = item.bbox / (img_w, img_h) \n",
    "        item.polygon = item.polygon / (img_w, img_h)\n",
    "        \n",
    "        detections.append(\n",
    "            fo.Detection(label=item.category, bounding_box=item.bbox.xywh, confidence=item.score)\n",
    "        )\n",
    "        segmentations.append(\n",
    "            fo.Polyline(label=item.category, points=[item.polygon.points], closed=True, filled=False)\n",
    "        )\n",
    "    sample[\"predictions\"] = fo.Detections(detections=detections)\n",
    "    sample[\"polylines\"] = fo.Polylines(polylines=segmentations)\n",
    "    samples.append(sample)\n",
    "    return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "d52d1001-5ab4-4dc7-a94f-33b322ec7f1b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "20\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 20/20 [00:05<00:00,  3.91it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "总处理时间: 5.68 秒\n",
      "总图片数量: 20\n",
      "FPS（帧率）: 3.52 张/秒\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "im_files =[]\n",
    "\n",
    "light = \"blue\"\n",
    "for batch in batch_nums:\n",
    "    im_files += glob.glob(f\"{osp.join(root_dir, batch)}/**/*{light}.png\", recursive=True)\n",
    "    \n",
    "print(len(im_files))\n",
    "start_time = time.time()  \n",
    "with ThreadPoolExecutor(max_workers=12) as executor:  # 设置最大线程数\n",
    "    results = list(tqdm(executor.map(process_image, im_files), total=len(im_files)))\n",
    "\n",
    "end_time = time.time()  # 记录结束时间\n",
    "total_time = end_time - start_time  # 总处理时间（秒）\n",
    "\n",
    "num_processed = len(im_files)  # 过滤掉 None 或 0\n",
    "fps = num_processed / total_time if total_time > 0 else 0\n",
    "\n",
    "print(f\"总处理时间: {total_time:.2f} 秒\")\n",
    "print(f\"总图片数量: {num_processed}\")\n",
    "print(f\"FPS（帧率）: {fps:.2f} 张/秒\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "20809523-899c-49df-8323-09c8e1dc2326",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 100% |███████████████████| 20/20 [138.5ms elapsed, 0s remaining, 144.4 samples/s] \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"800\"\n",
       "            src=\"http://0.0.0.0:8989/?notebook=True&subscription=bb1b8a5a-c1ae-406e-9efa-e9923349deab\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x703926fd3740>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataset.add_samples(samples)\n",
    "session = fo.launch_app(dataset, address=\"0.0.0.0\", port=8989)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6d563cd-f1f2-4594-9790-52167b936f18",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "147717b9-4adc-4ecf-92c8-24f90f89f190",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
